# 74/100 堆-数组中的第K个最大元素
# leetcode第215题: https://leetcode.cn/problems/kth-largest-element-in-an-array/?envType=study-plan-v2&envId=top-100-liked
# Date: 2024/11/28
from random import random


def findKthLargest(nums: list[int], k: int) -> int:
    nums.sort(reverse=True)
    return nums[k - 1]


def findKthLargest_quicksort(nums: list[int], k: int) -> int:
    """使用快速排序"""
    if len(nums) < 2:
        return nums[0]
    if k == 1:
        return max(nums)

    def quick_select(nums, k):
        # TODO 快排还要在实现一次
        # 随机选择基准数
        pivot = random.choice(nums)
        big, equal, small = [], [], []
        # 将大于、小于、等于 pivot 的元素划分至 big, small, equal 中
        for num in nums:
            if num > pivot:
                big.append(num)
            elif num < pivot:
                small.append(num)
            else:
                equal.append(num)
        if k <= len(big):
            # 第 k 大元素在 big 中，递归划分
            return quick_select(big, k)
        if len(nums) - len(small) < k:
            # 第 k 大元素在 small 中，递归划分
            return quick_select(small, k - len(nums) + len(small))
        # 第 k 大元素在 equal 中，直接返回 pivot
        return pivot

    return quick_select(nums, k)


if __name__ == '__main__':
    print(findKthLargest([3, 2, 1, 5, 6, 4], 2))  # 5
    print(findKthLargest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4))  # 4
